import os
from pathlib import Path
import rclpy
from rclpy.node import Node
from rclpy.time import Time
from rclpy.duration import Duration
import tf2_ros
from semantic_nav.json_utils import process_time_json, write_json
from semantic_nav.bag_utils import BagUtilsBase, BagTimeUtils
from semantic_nav.geomety_utils import pose_to_transform
from semantic_nav.log_utils import get_logger
logger = get_logger()
# Replace these with your actual frames, and topics
source_frame = "base_link"
target_frame = "map"
pose_topic = "/pose_topic"
image_topic = "/camera/color/image_raw"


class TfBagProcessor(BagUtilsBase, Node):
    def __init__(self, bag_path):
        super().__init__(bag_path)
        super(BagUtilsBase, self).__init__('tf_bag_processor')
        self.tf_buffer = tf2_ros.Buffer(cache_time=Duration(seconds=1e9))
        self.tf_listener = tf2_ros.TransformListener(self.tf_buffer, self)
        self.bag_path = bag_path
        self.first_timestamp = None

    def callback(self, topic, msg, timestamp):
        if topic == pose_topic:
            transform = pose_to_transform(msg, source_frame)
            self.tf_buffer.set_transform(transform, 'default_authority')
        elif topic in ['/tf', '/tf_static']:
            for transform in msg.transforms:
                self.tf_buffer.set_transform(transform, 'default_authority')
        return True

    def process_bag(self):
        bag_time_utils = BagTimeUtils(self.bag_path)
        self.first_timestamp = bag_time_utils.get_first_timestamp(image_topic)
        logger.info(self.first_timestamp)
        self.process()
        logger.info("Finished processing bag file and storing transforms in tf_buffer")

    def query_transforms(self, time_json_path):
        map_dict = process_time_json(time_json_path)

        if self.first_timestamp is None:
            raise ValueError("No images found in bag file.")
        logger.info(f'first timestamp: {self.first_timestamp}')
        for scene in map_dict:
            for obj in map_dict[scene]:
                logger.info(f"\nobject name: {obj}")
                timestamp = self.first_timestamp + Duration(seconds=map_dict[scene][obj])
                try:
                    transform = self.tf_buffer.lookup_transform(target_frame, source_frame, timestamp)
                    map_dict[scene][obj] = {
                        "position": {
                            "x": transform.transform.translation.x,
                            "y": transform.transform.translation.y,
                            "z": transform.transform.translation.z
                        },
                        "orientation": {
                            "x": transform.transform.rotation.x,
                            "y": transform.transform.rotation.y,
                            "z": transform.transform.rotation.z,
                            "w": transform.transform.rotation.w
                        }
                    }
                except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException) as e:
                    logger.error(f"Error looking up transform: {e}")
        output_path = Path(time_json_path).parent
        write_json(map_dict, str(output_path.joinpath('map.json')))


def main(args=None):
    rclpy.init(args=args)
    from ament_index_python.packages import get_package_prefix
    package_path = get_package_prefix('semantic_nav')
    bag_path = os.path.join(package_path, "bag", "test_bag")
    bag_processor = TfBagProcessor(bag_path)
    bag_processor.process_bag()
    time_json_path = os.path.join(package_path, "json/time.json")
    bag_processor.query_transforms(time_json_path)

    rclpy.shutdown()


if __name__ == "__main__":
    main()
